from typing import TypeAlias

import gymnasium as gym
import numpy as np
import polars as pl
from numpy.typing import NDArray

from .config import HomeEnvConfig

FloatArray: TypeAlias = NDArray[np.float32]

__all__ = ["HomeEnv"]


class HomeEnv(gym.Env):
    data_scope: FloatArray
    solar_data_table: pl.DataFrame
    load_data_table: pl.DataFrame
    temp_outdoor_data_table: pl.DataFrame
    price_data_table: pl.DataFrame
    ev_available_table: pl.DataFrame
    ev_reduction_table: pl.DataFrame
    ev_level_bkp: float
    state: FloatArray
    hour: int
    day: int
    n_step: int

    def __init__(self, **kwargs) -> None:
        self.config = HomeEnvConfig.from_dict(kwargs)
        # 加载数据表
        self.read_table()
        # 获取数据范围供归一化使用
        self.data_scope = self.get_data_scope()
        # 设置动作空间和观测空间
        if self.config.ev.enable:
            ev_action_low_bound = (-self.config.ev.max_p if self.config.ev.v2g else 0,)
            ev_action_high_bound = (self.config.ev.max_p,)
        else:
            ev_action_low_bound = ()
            ev_action_high_bound = ()
        self.action_space = gym.spaces.Box(
            low=np.array(
                [-self.config.ess.max_p, 0, *ev_action_low_bound], dtype=np.float32
            ),
            high=np.array(
                [self.config.ess.max_p, self.config.hvac.max_p, *ev_action_high_bound],
                dtype=np.float32,
            ),
        )
        if self.config.ev.enable:
            ev_state_bound = (
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
            )
        else:
            ev_state_bound = ()
        obs_low_bound = -np.array(
            [
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
                *ev_state_bound,
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
                np.finfo(np.float32).max,
            ]
        )
        self.observation_space = gym.spaces.Box(
            low=obs_low_bound, high=-obs_low_bound, dtype=np.float32
        )

    def read_table(self) -> None:
        deadline = self.config.duration.test + self.config.duration.train + 1
        columns = ["{}:00".format(i) for i in range(24)]
        schema = dict(zip(columns, [pl.Float32] * len(columns)))
        self.solar_data_table = pl.read_csv(
            self.config.external.solar,
            columns=columns,
            dtypes=schema,
            n_rows=deadline,
        )
        self.load_data_table = pl.read_csv(
            self.config.external.load,
            columns=columns,
            dtypes=schema,
            n_rows=deadline,
        )
        self.temp_outdoor_data_table = pl.read_csv(
            self.config.external.temp_outdoor,
            columns=columns,
            dtypes=schema,
            n_rows=deadline,
        )
        self.price_data_table = pl.read_csv(
            self.config.external.price,
            columns=columns,
            dtypes=schema,
            n_rows=deadline,
        )
        if self.config.ev.enable:
            self.ev_available_table = pl.read_csv(
                self.config.external.ev_available,
                columns=columns,
                dtypes=schema,
                n_rows=deadline,
            )
            self.ev_reduction_table = pl.read_csv(
                self.config.external.ev_reduction,
                n_rows=deadline,
            )

    def get_data_scope(self) -> FloatArray:
        """获取数据最大最小值, 供归一化使用"""
        if self.config.ev.enable:
            ev_low_bound = (self.config.ev.min_level, 0, 0)
            ev_high_bound = (self.config.ev.max_level, 1, 23)
        else:
            ev_low_bound = ()
            ev_high_bound = ()
        return np.array(
            [
                [
                    self.solar_data_table.min().min_horizontal().item(),
                    self.load_data_table.min().min_horizontal().item(),
                    self.config.ess.min_level,
                    *ev_low_bound,
                    self.temp_outdoor_data_table.min().min_horizontal().item(),
                    self.config.hvac.min_t,
                    self.price_data_table.min().min_horizontal().item(),
                    0,
                ],
                [
                    self.solar_data_table.max().max_horizontal().item(),
                    self.load_data_table.max().max_horizontal().item(),
                    self.config.ess.max_level,
                    *ev_high_bound,
                    self.temp_outdoor_data_table.max().max_horizontal().item(),
                    self.config.hvac.max_t,
                    self.price_data_table.max().max_horizontal().item(),
                    23,
                ],
            ],
            dtype=np.float32,
        )

    def normalize_state(self, s: FloatArray) -> FloatArray:
        """归一化状态"""
        assert len(s.shape) == 1
        return (s - self.data_scope[0]) / (self.data_scope[1] - self.data_scope[0])

    def denormalize_state(self, s: FloatArray) -> FloatArray:
        """恢复被归一化的状态"""
        assert len(s.shape) == 1
        return s * (self.data_scope[1] - self.data_scope[0]) + self.data_scope[0]

    def clip_hvac_action(self, p: float) -> float:
        idx = 7 if self.config.ev.enable else 4
        temp_indoor = self.state[idx]
        if temp_indoor <= self.config.hvac.min_t:
            p_hvac = 0
        elif temp_indoor > self.config.hvac.max_t:
            p_hvac = np.clip(p, 0.1, self.config.hvac.max_p)
        else:
            p_hvac = p
        return p_hvac

    def clip_ess_action(self, p: float) -> float:
        ess_level = self.state[2]
        if p >= 0:
            p_ess = np.clip(
                p,
                0,
                min(
                    (self.config.ess.max_level - ess_level) / self.config.ess.eta,
                    self.config.ess.max_p,
                ),
            )
        else:
            p_ess = -np.clip(
                -p,
                0,
                min(
                    (ess_level - self.config.ess.min_level) * self.config.ess.eta,
                    self.config.ess.max_p,
                ),
            )
        return p_ess

    def clip_ev_action(self, p: float) -> float:
        ev_ok, ev_level = self.state[4], self.state[3]
        if ev_ok == 0:
            p_ev = 0
        elif p > 0:
            p_ev = np.clip(
                p,
                0,
                min(
                    (self.config.ev.max_level - ev_level) / self.config.ev.eta,
                    self.config.ev.max_p,
                ),
            )
        else:
            p_ev = -np.clip(
                -p,
                0,
                min(
                    (ev_level - self.config.ev.min_level) * self.config.ev.eta,
                    self.config.ev.max_p,
                ),
            )
        return p_ev

    def get_ev_departure(self) -> int:
        """获取ev离开时间"""
        day, hour = self.day, self.hour
        cursor = 0
        if self.ev_available_table.item(day, hour) == 0:
            return -1
        while self.ev_available_table.item(day, hour) == 1:
            cursor += 1
            hour += 1
            if hour > 23:
                hour = 0
                day += 1
        return (day - self.day) * 24 + hour - self.hour

    def get_carbon_cost(self, p_grid: float) -> float:
        """计算碳交易成本"""
        cost = self.config.carbon_emission.price * (
            p_grid * self.config.carbon_emission.coef
            - self.config.carbon_emission.quota
        )
        return max(cost, 0)

    def step(self, a: FloatArray) -> tuple[FloatArray, float, bool, bool, dict]:
        """执行环境转移动态进入下一状态，并计算奖励"""
        """裁剪不合理的动作"""
        if self.config.ev.enable:
            p_ess, p_hvac, p_ev = a
            p_ev = self.clip_ev_action(p_ev)
        else:
            p_ess, p_hvac = a
        p_ess = self.clip_ess_action(p_ess)
        p_hvac = self.clip_hvac_action(p_hvac)
        if self.config.ev.enable:
            (
                p_solar,
                p_load,
                ess_level,
                ev_level,
                ev_ok,
                ev_departure,
                temp_outdoor,
                temp_indoor,
                price,
                _,
            ) = self.state
        else:
            p_solar, p_load, ess_level, temp_outdoor, temp_indoor, price, _ = self.state

        """ESS状态更新"""
        ess_level_next = ess_level + (
            p_ess * self.config.ess.eta if p_ess > 0 else p_ess / self.config.ess.eta
        )
        ess_level_next = np.clip(
            ess_level_next, self.config.ess.min_level, self.config.ess.max_level
        )

        """室温状态更新"""
        temp_indoor_next = self.config.hvac.epsilon * temp_indoor + (
            1 - self.config.hvac.epsilon
        ) * (temp_outdoor - self.config.hvac.eta * p_hvac / self.config.hvac.a)

        """EV状态更新"""
        if self.config.ev.enable:
            if ev_ok != 0:
                ev_level_next = ev_level + (
                    p_ev * self.config.ev.eta if p_ev > 0 else p_ev / self.config.ev.eta
                )
                ev_level_next = np.clip(
                    ev_level_next, self.config.ev.min_level, self.config.ev.max_level
                )
            else:
                ev_level_next = -1
        else:
            p_ev = 0

        """时间状态更新"""
        self.n_step += 1
        self.hour += 1
        if self.hour > 23:
            self.hour = 0
            self.day += 1

        if self.config.ev.enable:
            # 检测到0-1上升沿时重新加载soc, 根据离家时电量进行衰减
            ev_departure_next = self.get_ev_departure()
            ev_ok_next = ev_departure_next > 0
            if ev_ok == 0 and ev_ok_next != 0:
                ev_level_next = max(
                    self.ev_level_bkp - self.ev_reduction_table.item(self.day, 0),
                    self.config.ev.min_level,
                )

        """计算reward"""
        p_grid = p_load + p_ess + p_hvac + p_ev - p_solar  # 计算向主电网的购电/售电量
        c1 = p_grid * price
        if p_grid < 0:
            c1 *= self.config.grid.price_ratio
        c1 += abs(p_ess) * self.config.ess.aging_coef
        if self.config.carbon_emission.enable:
            c1 += self.get_carbon_cost(p_grid)
        c2 = max(0, self.config.hvac.min_t - temp_indoor_next) + max(
            0, temp_indoor_next - self.config.hvac.max_t
        )
        r = -float(c1 + c2)
        if self.config.ev.enable:
            # 到达depature时刻，如果电池未充满，增加一个惩罚项
            c3 = 0
            if ev_ok != 0 and ev_ok_next == 0:
                c3 = self.config.ev.max_level - ev_level_next
                r -= c3
                self.ev_level_bkp = ev_level_next

        """更新内部状态，内部状态始终不会被归一化"""
        ev_state = ()
        if self.config.ev.enable:
            ev_state = (ev_level_next, ev_ok_next, ev_departure_next)
        self.state = np.array(
            [
                self.solar_data_table.item(self.day, self.hour),
                self.load_data_table.item(self.day, self.hour),
                ess_level_next,
                *ev_state,
                self.temp_outdoor_data_table.item(self.day, self.hour),
                temp_indoor_next,
                self.price_data_table.item(self.day, self.hour),
                self.hour,
            ],
            dtype=np.float32,
        )

        """归一化输出状态"""
        if self.config.auto_normalize is True:
            state = self.normalize_state(self.state).astype(np.float32)
        else:
            state = self.state.copy()

        """判断是否结束"""
        t1 = False
        t2 = self.n_step >= self.config.max_days * 24
        info = {
            "c1": float(c1),
            "c2": float(c2),
            "p_grid": float(p_grid),
            "a": np.array([p_ess, p_hvac]),
        }
        if self.config.ev.enable:
            info["c3"] = float(c3)
            info["a"] = np.append(info["a"], p_ev)
        if self.config.auto_normalize:
            info["s"] = self.state.copy()
        return state, r, t1, t2, info

    def reset(
        self, *, seed: int | None = None, options: dict | None = None
    ) -> tuple[FloatArray, dict]:
        """重置环境"""
        self.day = self.config.start_day
        self.n_step = 0
        self.hour = 0
        ev_state = ()
        if self.config.ev.enable:
            ev_state = (
                self.config.ev.init_level,
                self.ev_available_table.item(self.day, 0),
                self.get_ev_departure(),
            )
        self.state = np.array(
            [
                self.solar_data_table.item(self.day, 0),
                self.load_data_table.item(self.day, 0),
                self.config.ess.init_level,
                *ev_state,
                self.temp_outdoor_data_table.item(self.day, 0),
                self.config.hvac.init_indoor_temp,
                self.price_data_table.item(self.day, 0),
                0,
            ],
            dtype=np.float32,
        )
        info = {"day": self.day, "hour": self.hour}
        if self.config.auto_normalize is True:
            state = self.normalize_state(self.state)
            info["s"] = self.state.copy()
        else:
            state = self.state.copy()
        return state, info
